import threading
import socket
import sys
import select
import datetime
import time
import json
import Queue
import traceback

from lib import event
from lib import log

ACCEPT_NEW_CONNECTIONS_POLL = 0.5
RECV_BUFFER_SIZE = 4096

### Disconnect Message
DISCONNECT_MESSAGE="DISCONNECT"

### Heartbeat Constants
HEARTBEAT_MESSAGE="HEARTBEAT"
HEARTBEATACK_MESSAGE="HEARTBEATACK"
MAX_HEARTBEAT_ACK_WAIT = 20.0
HEARTBEAT_INTERVAL = 60.0

##
## NetworkEvent classes
##

class NetworkEvent(event.Event): pass

class NetworkConnectionEvent(NetworkEvent):
    def __init__(self, connected, **kwds):
        super(NetworkConnectionEvent, self).__init__('connection', connected, **kwds)

class NetworkReceivedDataEvent(NetworkEvent):
    def __init__(self, data, **kwds):
        super(NetworkReceivedDataEvent, self).__init__('received', data, **kwds)

##
## Network Client, Server and Connection classes
##

class NetworkServerConnection(threading.Thread):
    """
    NetworkServer also provides several methods and properties that another process
    can use to control it.
        sendQueue - Queue.Queue instance. Post dictionary messages which will be sent to the client
        stop - stops the thread.
    """
    def __init__(self, server, connectionId, connection, address, eventQueue):
        """
        NetworkServerConnection(server, connectionId, connection, address, eventQueue)
            server - NetworkServer instance
            connectionId - unique identifier for this connection
            connection - socket.socket instance
            address - (address, port) tuple
            eventQueue - EventQueue instance. This is where NetworkEvent events will be posted
        """
        super(NetworkServerConnection, self).__init__()
        self._server = server
        self._connectionId = connectionId
        self._connection = connection
        self._client_address = address
        self._eventQueue = eventQueue
        self._stop = threading.Event()
        self._sendQueue = Queue.Queue()
        ### Initialize heartbeat information
        self._lastHeartbeatSent = datetime.datetime.now()
        self._lastReceivedHeartbeatAcknowledgement = datetime.datetime.now()
        self._lastHeartbeatAcknowledged = True
        log.info("Connection created for {0}".format(self))
    
    @property
    def eventQueue(self): return self._eventQueue
    @property
    def sendQueue(self): return self._sendQueue
    @property
    def connection(self): return self._connection
    @property
    def connectionId(self): return self._connectionId
    @property
    def client_address(self): return self._client_address
    @property
    def server(self): return self._server
    @property
    def lastHeartbeatSent(self): return self._lastHeartbeatSent
    @property
    def lastReceivedHeartbeatAcknowledgement(self): return self._lastReceivedHeartbeatAcknowledgement
    @property
    def lastHeartbeatAcknowledged(self): return self._lastHeartbeatAcknowledged
    
    def __str__(self): return "{0}:{1} ({2})".format(
        self.client_address[0],
        self.client_address[1],
        self.connectionId)
    
    def stop(self):
        """
        Closes the client connection.
        """
        self._stop.set()

    def isSocketAlive(self):
        """
        Returns True if the socket is still alive / connected. False otherwise.
        """
        return not (not self.lastHeartbeatAcknowledged and (datetime.datetime.now() - self.lastHeartbeatSent > datetime.timedelta(seconds=MAX_HEARTBEAT_ACK_WAIT)))
    
    def heartBeatDue(self):
        """
        Returns True if the server should send a heartbeat to the client
        """
        return self.isSocketAlive() and self.lastHeartbeatAcknowledged and (datetime.datetime.now() - self.lastHeartbeatSent > datetime.timedelta(seconds=HEARTBEAT_INTERVAL))
    
    def sendHeartbeatToClient(self):
        """
        Sends a heartbeat message to client to verify connection
        """
        #log.info("sending heartbeat to client {0}".format(self))
        self.sendMessage({'type': HEARTBEAT_MESSAGE})
        self._lastHeartbeatSent = datetime.datetime.now()
        self._lastHeartbeatAcknowledged = False
    
    def sendDisconnectToClient(self):
        """
        Sends a disconnect message to client
        """
        self.sendMessage({'type': DISCONNECT_MESSAGE})

    def sendMessage(self, message):
        try:
            self.connection.send(json.dumps(message))
        except:
            log.error(traceback.format_exc())

    def receivedMessage(self, message):
        """
        Parses and handles message from client.
        """
        try:
            parsedMessage = json.loads(message)
            if self.isSocketAlive() and not self.lastHeartbeatAcknowledged and parsedMessage['type'] == HEARTBEATACK_MESSAGE:
                #log.info("heartbeat acknowledged by client {0}".format(self))
                self._lastHeartbeatAcknowledged = True
            else:
                # create network receive event and post it
                receive_event = NetworkReceivedDataEvent(parsedMessage, eventContext=self)
                self.eventQueue.registerEvent(receive_event)
        except:
            traceback.print_exc()
    
    def run(self):
        """
        Thread run loop for the client connection.
        """
        connection_event = NetworkConnectionEvent(True, eventContext = self)
        self.eventQueue.registerEvent(connection_event)
        while not self._stop.is_set():
            try:
                if not self.isSocketAlive():
                    log.info("socket with client {0} has closed".format(self))
                    break
                if self.heartBeatDue():
                    self.sendHeartbeatToClient()
                ready_to_read, ready_to_write, in_error = select.select(
                        [self.connection],
                        [self.connection],
                        [self.connection],
                        0)
                if ready_to_read:
                    data = ""
                    try:
                        data = self.connection.recv(RECV_BUFFER_SIZE)
                    except:
                        log.error(traceback.format_exc())
                    if len(data) > 0:
                        self.receivedMessage(data)
                if ready_to_write:
                    if not self.sendQueue.empty():
                        next_send = self.sendQueue.get_nowait()
                        self.sendMessage(next_send)
                if in_error:
                    print "in error?"
                time.sleep(0.1)
            except Exception as e:
                traceback.print_exc()
                break
        log.info("Closing connection with {0}".format(self.client_address))
        self.connection.close()
        self.server.connectionClosed(self)
        connection_event = NetworkConnectionEvent(False, eventContext = self)
        self.eventQueue.registerEvent(connection_event)
        
class NetworkServer(threading.Thread):
    """
    NetworkServer manages client connections.
    """
    def __init__(self, address, port, eventQueue):
        super(NetworkServer, self).__init__()
        self._address = address
        self._port = port
        self._eventQueue = eventQueue
        self._stop = threading.Event()
        self._socket = None
        self._connections = {}
        
    @property
    def address(self): return self._address
    @property
    def port(self): return self._port
    @property
    def server(self): return self._server
    @property
    def eventQueue(self): return self._eventQueue
    @property
    def socket(self): return self._socket
    @property
    def connections(self): return self._connections
    
    def stop(self):
        self._stop.set()
    
    def closeConnection(self, clientConnection):
        """
        closeConnection
        """
        clientConnection.sendDisconnectToClient()
        clientConnection.stop()
        
    def connectionClosed(self, clientConnection):
        log.info("Client connection {0} reports that it is closed.".format(clientConnection))
        del self.connections[clientConnection.connectionId]
    
    def run(self):
        assert self.socket is None, 'socket is not None!'
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.bind((self.address, self.port))
        self._socket.listen(5)
        self._socket.setblocking(0)
        log.info("Network server listening on {0}:{1}".format(self.address, self.port))
        while not self._stop.is_set():
            try:
                conn, addr = self._socket.accept()
                conn_id = "{0}:{1}".format(addr[0],addr[1])
                client_connection = NetworkServerConnection(self, conn_id, conn, addr, self.eventQueue)
                self.connections[client_connection.connectionId] = client_connection
                client_connection.start()
            except socket.error:
                time.sleep(ACCEPT_NEW_CONNECTIONS_POLL)
        log.info("Closing all connections")
        for client_id, client_connection in self.connections.iteritems():
            log.info("Closing {0}".format(client_id))
            client_connection.stop()
        self._client_connections = {}
        
class NetworkClient(threading.Thread):
    """
    Connects to NetworkServer and communicates with it.
    """
    def __init__(self, address, port, eventQueue):
        super(NetworkClient, self).__init__()
        self._address = address
        self._port = port
        self._eventQueue = eventQueue
        self._socket = None
        self._stop = threading.Event()
        self._sendQueue = Queue.Queue()
    
    @property
    def sendQueue(self): return self._sendQueue
    @property
    def socket(self): return self._socket
    @property
    def eventQueue(self): return self._eventQueue
    
    def heartbeatAcknowledge(self):
        """
        Acknowledges a heartbeat sent from the server
        """
        self.sendQueue.put_nowait({'type': HEARTBEATACK_MESSAGE})
    
    def stop(self):
        self._stop.set()
    
    def sendMessage(self, message):
        message_data = json.dumps(message)
        log.info("sending to server >> " + message_data)
        self.socket.send(message_data)
        
    def receivedMessage(self, message):
        """
        Parses and handles message from server.
        """
        try:
            parsedMessage = json.loads(message)
            receive_event = NetworkReceivedDataEvent(parsedMessage, eventContext = self)
            self._eventQueue.registerEvent(receive_event)
        except:
            traceback.print_exc()
        
    def run(self):
        assert self._socket is None, 'socket is not None!'
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._socket.connect((self._address, self._port))
        self.eventQueue.registerEvent(NetworkConnectionEvent(True, eventContext=self))
        while not self._stop.is_set():
            ready_to_read, ready_to_write, in_error = select.select(
                    [self._socket],
                    [self._socket],
                    [self._socket])
            if ready_to_read:
                data = self._socket.recv(RECV_BUFFER_SIZE)
                if len(data) > 0:
                    self.receivedMessage(data)
            if ready_to_write:
                if not self.sendQueue.empty():
                    next_send = self.sendQueue.get_nowait()
                    self.sendMessage(next_send)
            time.sleep(0.1)     
        log.info("Closing connection with server")
        self._socket.close()
        self.eventQueue.registerEvent(NetworkConnectionEvent(False, eventContext=self))

##
## network client / server test harness
## below these lines
##
## To run the server:
## python -m lib.net server localhost:8000
##
## To run the client:
## python -m lib.net client localhost:8000
##

def test_server(address, port):
    import importlib
    event_module = importlib.import_module('lib.event')
    event_queue = event_module.EventQueue()
    server = NetworkServer(address, port, event_queue)
    log.info("Starting server")
    server.start()
    log.info("Server started -- press control-c to stop")
    while True:
        try:
            if not event_queue.noEvents:
                next_event = event_queue.nextEvent()
                if isinstance(next_event, NetworkReceivedDataEvent):
                    client_data = next_event.eventData
                    client_connection = next_event.eventContext
                if client_data['type'] == 'DISCONNECT':
                    server.closeConnection(client_connection)
                log.info("{0} => \"{1}\"".format(str(client_connection), client_data))
            time.sleep(3.0)
        except:
            print traceback.print_exc()
            break
    log.info("Stopping server") 
    server.stop()

def test_client(address, port):
    import importlib
    event_module = importlib.import_module('lib.event')
    event_queue = event_module.EventQueue()
    keyboard_module = importlib.import_module('lib.keyboard')
    keyboard_poller = keyboard_module.KeyboardPollerThread(event_queue)
    keyboard_poller.start()
    client = NetworkClient(address, port, event_queue)
    client.start()
    log.info("Enter 'quit' to quit.")
    while True:
        try:
            if not event_queue.noEvents:
                next_event = event_queue.nextEvent()
                if isinstance(next_event, keyboard_module.KeyboardEvent):
                    command = next_event.eventData
                    if command.lower() == 'quit':
                        break
                    else:
                        client.sendQueue.put_nowait({'type': 'command', 'command': command})
                elif isinstance(next_event, NetworkReceivedDataEvent):
                    message = next_event.eventData
                    if message['type'] == HEARTBEAT_MESSAGE:
                        client.heartbeatAcknowledge()
                    else:
                        print ">> ", next_event.eventDatas
                elif isinstance(next_event, NetworkConnectionEvent):
                    log.info("Connection event: {0}".format(next_event.eventData))
            time.sleep(0.01)
        except:
            break
    print "Breaking connection and stopping"
    client.sendMessage({'type': 'DISCONNECT'})
    client.stop()
    keyboard_poller.stop()

if __name__ == '__main__':
    test_mode = sys.argv[1].lower()
    test_address, test_port = sys.argv[2].split(":")
    if test_mode == 'server':
        test_server(test_address, int(test_port))
    elif test_mode == 'client':
        test_client(test_address, int(test_port))
    
